import pytest  # used for our unit tests
from workload import AlexNet

# unit tests

# ----------- Basic Test Cases -----------

def test_single_image_all_zeros():
    # Test with a single image of all zeros
    net = AlexNet()
    img = [0] * net.features_size
    codeflash_output = net._extract_features([img]); out = codeflash_output # 1.19μs -> 391ns (205% faster)

def test_single_image_all_ones():
    # Test with a single image of all ones
    net = AlexNet()
    img = [1] * net.features_size
    codeflash_output = net._extract_features([img]); out = codeflash_output # 1.13μs -> 351ns (222% faster)

def test_multiple_images_varied_values():
    # Test with multiple images with varied values
    net = AlexNet()
    img1 = [1] * net.features_size
    img2 = [2] * net.features_size
    img3 = list(range(net.features_size))
    codeflash_output = net._extract_features([img1, img2, img3]); out = codeflash_output # 1.01μs -> 381ns (165% faster)

def test_single_image_negative_values():
    # Test with negative values
    net = AlexNet()
    img = [-1] * net.features_size
    codeflash_output = net._extract_features([img]); out = codeflash_output # 1.18μs -> 381ns (210% faster)

def test_single_image_floats():
    # Test with floating point values
    net = AlexNet()
    img = [0.5] * net.features_size
    codeflash_output = net._extract_features([img]); out = codeflash_output # 1.13μs -> 351ns (223% faster)

# ----------- Edge Test Cases -----------

def test_empty_input_list():
    # Test with empty input list
    net = AlexNet()
    codeflash_output = net._extract_features([]); out = codeflash_output # 892ns -> 330ns (170% faster)





def test_image_with_nan_inf():
    # Test with image containing float('nan') and float('inf')
    net = AlexNet()
    img_nan = [float('nan')] * net.features_size
    img_inf = [float('inf')] * net.features_size
    # Should not raise, but mean/max/min/sum will be nan or inf
    codeflash_output = net._extract_features([img_nan]); out_nan = codeflash_output # 1.47μs -> 451ns (226% faster)
    codeflash_output = net._extract_features([img_inf]); out_inf = codeflash_output # 461ns -> 150ns (207% faster)


def test_large_batch_size():
    # Test with a large batch of images (e.g., 500)
    net = AlexNet()
    img = [1] * net.features_size
    batch = [img] * 500
    codeflash_output = net._extract_features(batch); out = codeflash_output # 8.44μs -> 451ns (1771% faster)
    for o in out:
        pass

def test_large_image_values():
    # Test with large values in image
    net = AlexNet()
    img = [1e6] * net.features_size
    codeflash_output = net._extract_features([img]); out = codeflash_output # 1.32μs -> 370ns (257% faster)

def test_large_varied_images():
    # Test with 100 varied images
    net = AlexNet()
    batch = []
    for i in range(100):
        img = [i] * net.features_size
        batch.append(img)
    codeflash_output = net._extract_features(batch); out = codeflash_output
    for i, o in enumerate(out):
        pass

def test_performance_large_batch():
    # Test with a batch of 999 images to check performance/scalability
    net = AlexNet()
    img = [2] * net.features_size
    batch = [img] * 999
    codeflash_output = net._extract_features(batch); out = codeflash_output # 16.3μs -> 391ns (4061% faster)
    for o in out:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.